'''
 # @ Author: Justin
 # @ Create Time: 2025-02-11 21:32:14
 # @ Modified by: Justin
 # @ Modified time: 2025-02-11 21:32:17
 # @ Description: 去除deepseek输出的<think>标签、markdown格式的特殊字符以及LaTeX公式
 '''

import markdown
from bs4 import BeautifulSoup
import re

def parse_latex(text):
    '''
    去除LaTeX公式的特殊字符
    Args:
        text: 输入文本
    Returns:
        去除LaTeX公式的特殊字符后的文本
    '''

    text = re.sub(r'\\frac\{([^}]+)\}\{([^}]+)\}', r'\1/\2', text)  # 分数
    text = re.sub(r'\\sqrt\{([^}]+)\}', r'sqrt(\1)', text)         # 平方根
    text = re.sub(r'\\Delta', 'delta', text)                       # Delta
    text = re.sub(r'\\neq', '≠', text)                             # 不等于
    text = re.sub(r'\\pm', '±', text)                              # 加减号
    text = re.sub(r'\\times', '×', text)                           # 乘号
    text = re.sub(r'\\div', '÷', text)                             # 除号
    text = re.sub(r'\\in', '∈', text)                              # 属于
    text = re.sub(r'\\forall', '∀', text)                          # 全称量词
    text = re.sub(r'\\exists', '∃', text)                          # 存在量词
    text = re.sub(r'\\rightarrow', '→', text)                      # 右箭头
    text = re.sub(r'\\leftarrow', '←', text)                       # 左箭头
    text = re.sub(r'\\Rightarrow', '⇒', text)                      # 双箭头
    text = re.sub(r'\\Leftarrow', '⇐', text)                       # 双左箭头
    text = re.sub(r'\\leftrightarrow', '↔', text)                  # 左右箭头
    text = re.sub(r'\\iff', '⇔', text)                             # 当且仅当
    text = re.sub(r'\\geq', '≥', text)                             # 大于等于
    text = re.sub(r'\\leq', '≤', text)                             # 小于等于
    text = re.sub(r'\\approx', '≈', text)                          # 约等于
    text = re.sub(r'\\sim', '∼', text)                             # 相似
    text = re.sub(r'\\equiv', '≡', text)                           # 等价于
    text = re.sub(r'\\cdot', '·', text)                            # 点乘
    text = re.sub(r'\\circ', '∘', text)                            # 圆圈运算符
    text = re.sub(r'\\cup', '∪', text)                             # 并集
    text = re.sub(r'\\cap', '∩', text)                             # 交集
    text = re.sub(r'\\subset', '⊂', text)                          # 子集
    text = re.sub(r'\\supset', '⊃', text)                          # 超集
    text = re.sub(r'\\subseteq', '⊆', text)                        # 子集或等于
    text = re.sub(r'\\supseteq', '⊇', text)                        # 超集或等于
    text = re.sub(r'\\emptyset', '∅', text)                        # 空集
    text = re.sub(r'\\infty', '∞', text)                           # 无穷大
    text = re.sub(r'\\angle', '∠', text)                           # 角度
    text = re.sub(r'\\perp', '⊥', text)                            # 垂直
    text = re.sub(r'\\parallel', '∥', text)                        # 平行
    text = re.sub(r'\\lceil', '⌈', text)                           # 左上取整
    text = re.sub(r'\\rceil', '⌉', text)                           # 右上取整
    text = re.sub(r'\\lfloor', '⌊', text)                          # 左下取整
    text = re.sub(r'\\rfloor', '⌋', text)                          # 右下取整
    text = re.sub(r'\\left', '', text)                             # 替换 \left
    text = re.sub(r'\\right', '', text)                            # 替换 \right
    text = re.sub(r'\\left\(', '(', text)                          # 左括号
    text = re.sub(r'\\right\)', ')', text)                         # 右括号
    text = re.sub(r'\\left\[', '[', text)                          # 左中括号
    text = re.sub(r'\\right\]', ']', text)                         # 右中括号
    text = re.sub(r'\\left\{', '{', text)                          # 左花括号
    text = re.sub(r'\\right\}', '}', text)                         # 右花括号
    text = re.sub(r'\\left\|', '|', text)                          # 左竖线
    text = re.sub(r'\\right\|', '|', text)                         # 右竖线
    text = re.sub(r'\\left\<', '<', text)                          # 左尖括号
    text = re.sub(r'\\right\>', '>', text)                         # 右尖括号
    text = re.sub(r'\\left\=', '=', text)                          # 左等号
    text = re.sub(r'\\right\=', '=', text)                         # 右等号
    text = re.sub(r'\\left\.', '', text)                           # 左点
    text = re.sub(r'\\right\.', '', text)                          # 右点
    text = re.sub(r'\\left\,', ',', text)                          # 左逗号
    text = re.sub(r'\\right\,', ',', text)                         # 右逗号
    text = re.sub(r'\\left\;', ';', text)                          # 左分号
    text = re.sub(r'\\right\;', ';', text)                         # 右分号
    text = re.sub(r'\\left\:', ':', text)                          # 左冒号
    text = re.sub(r'\\right\:', ':', text)                         # 右冒号
    text = re.sub(r'\\left\?', '?', text)                          # 左问号
    text = re.sub(r'\\right\?', '?', text)                         # 右问号
    text = re.sub(r'\\left\!', '!', text)                          # 左感叹号
    text = re.sub(r'\\right\!', '!', text)                         # 右感叹号
    text = re.sub(r'\\left\^', '^', text)                          # 左上标
    text = re.sub(r'\\right\^', '^', text)                         # 右上标
    text = re.sub(r'\\left\_', '_', text)                          # 左下划线
    text = re.sub(r'\\right\_', '_', text)                         # 右下划线
    text = re.sub(r'\\left\~', '~', text)                          # 左波浪线
    text = re.sub(r'\\right\~', '~', text)                         # 右波浪线
    text = re.sub(r'\\left\`', '`', text)                          # 左反引号
    text = re.sub(r'\\right\`', '`', text)                         # 右反引号
    text = re.sub(r'\\left\'', '\'', text)                         # 左单引号
    text = re.sub(r'\\right\'', '\'', text)                        # 右单引号
    text = re.sub(r'\\left\"', '"', text)                          # 左双引号
    text = re.sub(r'\\right\"', '"', text)                         # 右双引号
    text = re.sub(r'\\left\#', '#', text)                          # 左井号
    text = re.sub(r'\\right\#', '#', text)                         # 右井号
    text = re.sub(r'\\left\$', '$', text)                          # 左美元符号
    text = re.sub(r'\\right\$', '$', text)                         # 右美元符号
    text = re.sub(r'\\left\%', '%', text)                          # 左百分号
    text = re.sub(r'\\right\%', '%', text)                         # 右百分号
    text = re.sub(r'\\left\&', '&', text)                          # 左和号
    text = re.sub(r'\\right\&', '&', text)                         # 右和号

    # 处理 \text{} 和 \boxed{}
    text = re.sub(r'\\text\{([^}]+)\}', r'\1', text)               # \text{}
    text = re.sub(r'\\boxed\{([^}]+)\}', r'\1', text)              # \boxed{}

    # 处理 \begin{} 和 \end{}
    text = re.sub(r'\\begin\{([^}]+)\}', r'\1', text)             # \begin{}
    text = re.sub(r'\\end\{([^}]+)\}', r'\1', text)               # \end{}

    # 其他
    text = re.sub(r'\\frac\{([^}]+)\}\{([^}]+)\}', r'\1/\2', text)    # 分数

    text = re.sub(r"\\quad", '  ', text)                            # tab

    text = re.sub(r"\\delta", 'Δ', text)                            # delta

    return text

# 解析输出的markdown文档
def parse_markdown(content):
    '''
    解析markdown文档
    Args:
        content: markdown文档内容
    Returns:
        纯文本内容
    '''
    html = markdown.markdown(content, extensions=['mdx_math'])

    # 使用 BeautifulSoup 提取纯文本
    soup = BeautifulSoup(html, 'html.parser')

    # 找到所有的数学公式并替换为原始 LaTeX 代码
    for formula in soup.find_all('script', type='math/tex'):
        latex = formula.text.strip()
        formula.replace_with(f' {latex} ')
    for formula in soup.find_all('script', type='math/tex; mode=display'):
        latex = formula.text.strip()
        formula.replace_with(f' {latex} ')

    # 提取纯文本
    text = soup.get_text()

    # 去掉多余的空格和换行
    text = re.sub(r'\s+', ' ', text).strip()

    return text

def parse(content):
    '''
    解析deepseek输出的markdown文档
    Args:
        content: deepseek输出的markdown文档内容
    Returns:
        纯文本内容
    '''
    
    # think标签
    content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)

    content = parse_markdown(content)
    content = parse_latex(content)

    content = re.sub(r'\s+', ' ', content).strip()

    return content
